# -*- coding: utf-8 -*-

import xlrd, os

from app.helper.helper import data_path


class Excel:
    """
    Excel简单封装
    """

    def __init__(self, filename=None, file_contents=None):
        if file_contents:
            self.workbook = xlrd.open_workbook(file_contents=file_contents)
        else:
            file_path = os.path.join(data_path(), filename)
            self.workbook = xlrd.open_workbook(file_path)

    def get_sheet_data(self, sheet_name, key_num=0, start_value_num=1):
        """
        根据sheet名称获取数据
        :param sheet_name: sheet名称
        :param key_num: key所在行编号，一般第一行为key
        :param start_value_num: value开始的行数
        :return:
        """
        if key_num > start_value_num:
            raise Exception("key_num value must be less than start_value_num value")

        sheet_table = self.workbook.sheet_by_name(sheet_name)
        return self.get_table_data(sheet_table, key_num, start_value_num)

    def get_excel_data(self, key_num=0, start_value_num=1):
        """
        获取整个excel数据
        :param key_num: key所在行编号，一般第一行为key
        :param start_value_num: value开始的行数
        :return:
        """
        sheets = self.workbook.sheet_names()

        data = []
        for sheet_name in sheets:
            data.append(self.get_sheet_data(sheet_name, key_num, start_value_num))

        return data

    def get_table_data(self, sheet_table, key_num=0, start_value_num=1):
        """获取sheet里面所有数据"""

        row_num = sheet_table.nrows

        result = []
        data_key = []

        for i in range(row_num):
            keys = sheet_table.row_values(i)

            if i == key_num:
                data_key = keys
            elif i >= start_value_num:
                data = {}
                for k in range(len(keys)):
                    data[data_key[k]] = keys[k]
                result.append(data)

        return result
